# -*- coding: utf-8 -*-
"""
集成模型管理器
负责管理多个预测算法，协调预测结果，提供统一的预测接口
"""

import os
import json
import numpy as np
import pandas as pd
from datetime import datetime
from typing import Dict, List, Any, Optional
from collections import defaultdict

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

class EnsembleModelManager:
    """
    集成模型管理器
    管理多个预测算法，提供统一的训练、预测和结果融合接口
    """
    
    def __init__(self, lottery_type: str = 'ssq', config: Dict[str, Any] = None):
        """
        初始化集成模型管理器
        
        Args:
            lottery_type: 彩票类型 ('ssq' 或 'dlt')
            config: 配置参数
        """
        self.lottery_type = lottery_type
        self.config = config or {}
        self.algorithms = {}
        self.algorithm_weights = {}
        self.prediction_history = []
        self.performance_metrics = {}
        
        # 设置彩票参数
        if lottery_type == 'ssq':
            self.red_ball_range = (1, 33)
            self.blue_ball_range = (1, 16)
            self.red_ball_count = 6
            self.blue_ball_count = 1
        elif lottery_type == 'dlt':
            self.red_ball_range = (1, 35)
            self.blue_ball_range = (1, 12)
            self.red_ball_count = 5
            self.blue_ball_count = 2
        else:
            raise ValueError(f"不支持的彩票类型: {lottery_type}")
        
        # 初始化模式分析器
        self.pattern_analyzer = None
        self._initialize_pattern_analyzer()
        
        # 默认算法权重
        self.default_weights = {
            'LSTM-Sequence': 0.25,
            'Gradient-Boosting': 0.25,
            'Pattern-Memory': 0.25,
            'Transformer-Pattern': 0.15,
            'Zone-Weight': 0.10
        }
        
        self.algorithm_weights = self.default_weights.copy()
    
    def _initialize_pattern_analyzer(self):
        """初始化模式分析器"""
        try:
            from pattern_analysis.consecutive_tracker import ConsecutivePatternTracker
            self.pattern_analyzer = ConsecutivePatternTracker(self.lottery_type)
        except ImportError:
            print("⚠️ 模式分析器未安装，将使用简化版本")
            self.pattern_analyzer = None
    
    def register_algorithm(self, algorithm_name: str, algorithm_instance, weight: float = 0.2):
        """
        注册预测算法
        
        Args:
            algorithm_name: 算法名称
            algorithm_instance: 算法实例
            weight: 算法权重
        """
        self.algorithms[algorithm_name] = algorithm_instance
        self.algorithm_weights[algorithm_name] = weight
        print(f"✅ 算法已注册: {algorithm_name} (权重: {weight})")
    
    def auto_register_algorithms(self):
        """自动注册所有可用的算法"""
        registered_count = 0
        print(f" 开始自动注册算法（当前彩票类型: {self.lottery_type}）...")
        
        # 尝试注册序列LSTM算法
        try:
            print("   🔄 尝试注册序列LSTM算法...")
            from algorithms.sequence_lstm import SequenceAwareLSTM
            
            # 动态计算参数
            input_dim = self.config.get('input_dim', 70)
            hidden_dim = self.config.get('hidden_dim', 128)
            output_dim = self.red_ball_range[1]  # 红球最大号码
            sequence_length = self.red_ball_count  # 红球数量
            
            lstm_algo = SequenceAwareLSTM(
                input_dim=input_dim,
                hidden_dim=hidden_dim,
                output_dim=output_dim,
                sequence_length=sequence_length
            )
            self.register_algorithm('LSTM-Sequence', lstm_algo, 0.25)
            registered_count += 1
            print("   ✅ 序列LSTM算法注册成功")
        except ImportError as e:
            print(f"   ⚠️ 序列LSTM算法导入失败: {e}")
        except Exception as e:
            print(f"   ❌ 序列LSTM算法注册失败: {e}")
        
        # 尝试注册梯度提升算法
        try:
            print("   🔄 尝试注册梯度提升算法...")
            from algorithms.gradient_boosting import GradientBoostingPredictor
            gb_algo = GradientBoostingPredictor(self.lottery_type)
            self.register_algorithm('Gradient-Boosting', gb_algo, 0.25)
            registered_count += 1
            print("   ✅ 梯度提升算法注册成功")
        except ImportError as e:
            print(f"   ⚠️ 梯度提升算法导入失败: {e}")
        except Exception as e:
            print(f"   ❌ 梯度提升算法注册失败: {e}")
        
        # 尝试注册模式记忆网络
        try:
            print("   🔄 尝试注册模式记忆网络...")
            from algorithms.pattern_memory import PatternMemoryNetwork
            pm_algo = PatternMemoryNetwork(
                self.lottery_type,
                memory_size=self.config.get('memory_size', 1000)
            )
            self.register_algorithm('Pattern-Memory', pm_algo, 0.25)
            registered_count += 1
            print("   ✅ 模式记忆网络注册成功")
        except ImportError as e:
            print(f"   ⚠️ 模式记忆网络导入失败: {e}")
        except Exception as e:
            print(f"   ❌ 模式记忆网络注册失败: {e}")
        
        # 如果主要算法注册失败，尝试注册备用简化算法
        if registered_count == 0:
            print("   🔍 主要算法注册失败，尝试注册备用算法...")
            try:
                # 注册一个简化的随机算法作为备用
                from algorithms.base_algorithm import BaseAlgorithm
                
                class SimpleRandomPredictor(BaseAlgorithm):
                    def __init__(self, lottery_type):
                        super().__init__(lottery_type)
                        self.algorithm_name = "Simple-Random"
                    
                    def train(self, historical_data):
                        return {"status": "completed", "method": "random"}
                    
                    def predict(self, recent_data, num_predictions=1):
                        import random
                        predictions = []
                        for i in range(num_predictions):
                            if self.lottery_type == 'dlt':
                                red = sorted(random.sample(range(1, 36), 5))
                                blue = random.sample(range(1, 13), 2)
                            else:
                                red = sorted(random.sample(range(1, 34), 6))
                                blue = [random.randint(1, 16)]
                            
                            predictions.append({
                                'red_balls': red,
                                'blue_balls': blue
                            })
                        return predictions
                
                backup_algo = SimpleRandomPredictor(self.lottery_type)
                self.register_algorithm('Backup-Random', backup_algo, 1.0)
                registered_count += 1
                print("   ✅ 备用随机算法注册成功")
                
            except Exception as e:
                print(f"   ❌ 备用算法注册也失败: {e}")
        
        print(f" 自动注册完成，成功注册 {registered_count} 个算法")
        
        # 显示注册结果的详细信息
        if registered_count > 0:
            print(f"   📋 已注册算法列表: {list(self.algorithms.keys())}")
            print(f"   ⚖️ 算法权重: {self.algorithm_weights}")
        else:
            print("   ⚠️ 没有成功注册任何算法")
        
        return registered_count
    
    def train_all_algorithms(self, historical_data: pd.DataFrame, 
                           validation_split: float = 0.2) -> Dict[str, Any]:
        """
        训练所有注册的算法
        
        Args:
            historical_data: 历史开奖数据
            validation_split: 验证集比例
            
        Returns:
            训练结果摘要
        """
        if not self.algorithms:
            print("⚠️ 没有注册的算法，尝试自动注册...")
            self.auto_register_algorithms()
        
        print(f"🚀 开始训练 {len(self.algorithms)} 个算法...")
        training_results = {}
        
        # 分割训练和验证数据
        split_idx = int(len(historical_data) * (1 - validation_split))
        train_data = historical_data[:split_idx]
        val_data = historical_data[split_idx:]
        
        # 模式分析
        pattern_analysis = None
        if self.pattern_analyzer:
            print(" 执行模式分析...")
            pattern_analysis = self.pattern_analyzer.analyze_consecutive_patterns(train_data)
            zone_analysis = self.pattern_analyzer.analyze_zone_trends(train_data)
            pattern_analysis.update(zone_analysis)
        
        # 训练每个算法
        for name, algorithm in self.algorithms.items():
            print(f"\n🔧 训练 {name}...")
            start_time = datetime.now()
            
            try:
                # 训练算法
                training_result = algorithm.train(train_data)
                
                # 在验证集上评估
                val_predictions = self._evaluate_on_validation(
                    algorithm, val_data, pattern_analysis
                )
                
                end_time = datetime.now()
                duration = (end_time - start_time).total_seconds()
                
                training_results[name] = {
                    'status': '✅ 成功',
                    'duration': duration,
                    'training_result': training_result,
                    'validation_predictions': len(val_predictions) if val_predictions else 0
                }
                
                print(f"   ✅ {name} 训练完成 ({duration:.2f}s)")
                
            except Exception as e:
                training_results[name] = {
                    'status': '❌ 失败',
                    'error': str(e)
                }
                print(f"   ❌ {name} 训练失败: {e}")
        
        # 更新算法权重
        self._update_algorithm_weights(training_results)
        
        return {
            'training_results': training_results,
            'pattern_analysis': pattern_analysis,
            'trained_algorithms': len([r for r in training_results.values() 
                                     if r['status'] == '✅ 成功'])
        }
    
    def generate_ensemble_predictions(self, recent_data: np.ndarray, 
                                    num_predictions: int = 3,
                                    fusion_method: str = 'weighted_voting') -> List[Dict[str, Any]]:
        """
        生成集成预测结果
        
        Args:
            recent_data: 最近的开奖数据
            num_predictions: 预测组数
            fusion_method: 融合方法 ('weighted_voting', 'majority_voting', 'stacking')
            
        Returns:
            集成预测结果列表
        """
        if not self.algorithms:
            raise ValueError("没有可用的算法，请先注册算法")
        
        print(f"🎯 生成集成预测 ({num_predictions}组，融合方法：{fusion_method})")
        
        # 获取模式分析结果
        pattern_analysis = {}
        if self.pattern_analyzer and hasattr(recent_data, 'shape'):
            try:
                # 转换为DataFrame格式进行模式分析
                recent_df = self._convert_to_dataframe(recent_data)
                pattern_analysis = self.pattern_analyzer.analyze_consecutive_patterns(recent_df)
                zone_analysis = self.pattern_analyzer.analyze_zone_trends(recent_df)
                pattern_analysis.update(zone_analysis)
            except Exception as e:
                print(f"⚠️ 模式分析失败: {e}")
        
        # 从每个算法获取预测
        algorithm_predictions = {}
        for name, algorithm in self.algorithms.items():
            try:
                if hasattr(algorithm, 'predict_multiple_sets'):
                    preds = algorithm.predict_multiple_sets(
                        recent_data, pattern_analysis, num_sets=num_predictions
                    )
                elif hasattr(algorithm, 'predict'):
                    preds = algorithm.predict(recent_data, num_predictions=num_predictions)
                else:
                    print(f"⚠️ {name} 没有预测方法")
                    continue
                
                algorithm_predictions[name] = preds
                print(f"   ✅ {name}: {len(preds)}组预测")
                
            except Exception as e:
                print(f"   ❌ {name} 预测失败: {e}")
                algorithm_predictions[name] = []
        
        # 融合预测结果
        ensemble_predictions = self._fuse_predictions(
            algorithm_predictions, num_predictions, fusion_method
        )
        
        # 记录预测历史
        self.prediction_history.append({
            'timestamp': datetime.now().isoformat(),
            'algorithm_predictions': algorithm_predictions,
            'ensemble_predictions': ensemble_predictions,
            'fusion_method': fusion_method
        })
        
        return ensemble_predictions
    
    def _evaluate_on_validation(self, algorithm, val_data: pd.DataFrame, 
                              pattern_analysis: Dict = None) -> List:
        """在验证集上评估算法"""
        try:
            if len(val_data) < 5:
                return []
            
            # 使用前几期预测后几期
            recent_data = val_data.iloc[-5:].values  # 最近5期
            
            if hasattr(algorithm, 'predict_multiple_sets'):
                predictions = algorithm.predict_multiple_sets(
                    recent_data, pattern_analysis or {}, num_sets=3
                )
            elif hasattr(algorithm, 'predict'):
                predictions = algorithm.predict(recent_data, num_predictions=3)
            else:
                return []
            
            return predictions
            
        except Exception as e:
            print(f"验证评估失败: {e}")
            return []
    
    def _update_algorithm_weights(self, training_results: Dict[str, Any]):
        """根据训练结果更新算法权重"""
        successful_algorithms = [name for name, result in training_results.items() 
                               if result['status'] == '✅ 成功']
        
        if not successful_algorithms:
            return
        
        # 重新分配权重给成功的算法
        equal_weight = 1.0 / len(successful_algorithms)
        new_weights = {}
        
        for name in self.algorithms.keys():
            if name in successful_algorithms:
                new_weights[name] = equal_weight
            else:
                new_weights[name] = 0.0
        
        self.algorithm_weights = new_weights
        print(f"🔄 算法权重已更新: {len(successful_algorithms)} 个算法平均分配")
    
    def _fuse_predictions(self, algorithm_predictions: Dict[str, List], 
                         num_predictions: int, method: str) -> List[Dict[str, Any]]:
        """融合多个算法的预测结果"""
        if method == 'weighted_voting':
            return self._weighted_voting_fusion(algorithm_predictions, num_predictions)
        elif method == 'majority_voting':
            return self._majority_voting_fusion(algorithm_predictions, num_predictions)
        elif method == 'stacking':
            return self._stacking_fusion(algorithm_predictions, num_predictions)
        else:
            return self._simple_combination_fusion(algorithm_predictions, num_predictions)
    
    def _weighted_voting_fusion(self, algorithm_predictions: Dict[str, List], 
                              num_predictions: int) -> List[Dict[str, Any]]:
        """加权投票融合"""
        ensemble_predictions = []
        
        for pred_idx in range(num_predictions):
            # 收集所有算法的第pred_idx个预测
            red_candidates = defaultdict(float)
            blue_candidates = defaultdict(float)
            
            total_weight = 0
            for algo_name, predictions in algorithm_predictions.items():
                if pred_idx < len(predictions):
                    pred = predictions[pred_idx]
                    weight = self.algorithm_weights.get(algo_name, 0.1)
                    total_weight += weight
                    
                    # 处理红球
                    red_balls = pred.get('red', pred.get('red_balls', []))
                    for ball in red_balls:
                        red_candidates[ball] += weight
                    
                    # 处理蓝球
                    blue_balls = pred.get('blue', pred.get('blue_balls', []))
                    for ball in blue_balls:
                        blue_candidates[ball] += weight
            
            # 选择权重最高的号码
            if red_candidates:
                sorted_red = sorted(red_candidates.items(), key=lambda x: x[1], reverse=True)
                selected_red = [ball for ball, _ in sorted_red[:self.red_ball_count]]
                selected_red = sorted(selected_red)
            else:
                # 备用方案：随机选择
                import random
                selected_red = sorted(random.sample(
                    range(self.red_ball_range[0], self.red_ball_range[1] + 1),
                    self.red_ball_count
                ))
            
            if blue_candidates:
                sorted_blue = sorted(blue_candidates.items(), key=lambda x: x[1], reverse=True)
                selected_blue = [ball for ball, _ in sorted_blue[:self.blue_ball_count]]
            else:
                # 备用方案：随机选择
                import random
                selected_blue = random.sample(
                    range(self.blue_ball_range[0], self.blue_ball_range[1] + 1),
                    self.blue_ball_count
                )
            
            ensemble_predictions.append({
                'red': selected_red,
                'blue': selected_blue,
                'algorithm': 'Ensemble-WeightedVoting',
                'confidence': min(total_weight, 1.0),
                'contributing_algorithms': len(algorithm_predictions)
            })
        
        return ensemble_predictions
    
    def _majority_voting_fusion(self, algorithm_predictions: Dict[str, List], 
                              num_predictions: int) -> List[Dict[str, Any]]:
        """多数投票融合"""
        ensemble_predictions = []
        
        for pred_idx in range(num_predictions):
            red_votes = defaultdict(int)
            blue_votes = defaultdict(int)
            
            # 统计投票
            for algo_name, predictions in algorithm_predictions.items():
                if pred_idx < len(predictions):
                    pred = predictions[pred_idx]
                    
                    # 红球投票
                    red_balls = pred.get('red', pred.get('red_balls', []))
                    for ball in red_balls:
                        red_votes[ball] += 1
                    
                    # 蓝球投票
                    blue_balls = pred.get('blue', pred.get('blue_balls', []))
                    for ball in blue_balls:
                        blue_votes[ball] += 1
            
            # 选择得票最多的号码
            selected_red = sorted([ball for ball, _ in 
                                 sorted(red_votes.items(), key=lambda x: x[1], reverse=True)
                                 [:self.red_ball_count]])
            
            selected_blue = [ball for ball, _ in 
                           sorted(blue_votes.items(), key=lambda x: x[1], reverse=True)
                           [:self.blue_ball_count]]
            
            ensemble_predictions.append({
                'red': selected_red,
                'blue': selected_blue,
                'algorithm': 'Ensemble-MajorityVoting',
                'confidence': 0.7,
                'contributing_algorithms': len(algorithm_predictions)
            })
        
        return ensemble_predictions
    
    def _simple_combination_fusion(self, algorithm_predictions: Dict[str, List], 
                                 num_predictions: int) -> List[Dict[str, Any]]:
        """简单组合融合"""
        ensemble_predictions = []
        all_predictions = []
        
        # 收集所有预测
        for predictions in algorithm_predictions.values():
            all_predictions.extend(predictions)
        
        # 随机选择不同的预测组合
        import random
        selected_predictions = random.sample(
            all_predictions, 
            min(num_predictions, len(all_predictions))
        )
        
        for i, pred in enumerate(selected_predictions):
            ensemble_predictions.append({
                'red': pred.get('red', pred.get('red_balls', [])),
                'blue': pred.get('blue', pred.get('blue_balls', [])),
                'algorithm': f"Ensemble-Combination-{i+1}",
                'confidence': pred.get('confidence', 0.5),
                'source_algorithm': pred.get('algorithm', 'Unknown')
            })
        
        return ensemble_predictions
    
    def _stacking_fusion(self, algorithm_predictions: Dict[str, List], 
                        num_predictions: int) -> List[Dict[str, Any]]:
        """堆叠融合（简化版本）"""
        # 这里实现一个简化的堆叠融合
        return self._weighted_voting_fusion(algorithm_predictions, num_predictions)
    
    def _convert_to_dataframe(self, data: np.ndarray) -> pd.DataFrame:
        """将数组数据转换为DataFrame格式"""
        if self.lottery_type == 'ssq':
            columns = ['Red_1', 'Red_2', 'Red_3', 'Red_4', 'Red_5', 'Red_6', 'Blue_1']
        else:  # dlt
            columns = ['Red_1', 'Red_2', 'Red_3', 'Red_4', 'Red_5', 'Blue_1', 'Blue_2']
        
        # 确保数据维度正确
        if len(data.shape) == 1:
            data = data.reshape(1, -1)
        
        # 只取前面几列作为彩票号码
        lottery_data = data[:, :len(columns)] if data.shape[1] >= len(columns) else data
        
        try:
            import pandas as pd
            return pd.DataFrame(lottery_data, columns=columns[:lottery_data.shape[1]])
        except ImportError:
            # 如果pandas不可用，返回简化格式
            return {
                'data': lottery_data,
                'columns': columns[:lottery_data.shape[1]]
            }
    
    def get_algorithm_status(self) -> Dict[str, Any]:
        """获取算法状态信息"""
        return {
            'registered_algorithms': list(self.algorithms.keys()),
            'algorithm_weights': self.algorithm_weights,
            'lottery_type': self.lottery_type,
            'total_predictions': len(self.prediction_history),
            'last_prediction': self.prediction_history[-1]['timestamp'] if self.prediction_history else None
        }
    
    def save_state(self, filepath: str) -> bool:
        """保存管理器状态"""
        try:
            state = {
                'lottery_type': self.lottery_type,
                'config': self.config,
                'algorithm_weights': self.algorithm_weights,
                'prediction_history': self.prediction_history[-100:],  # 只保存最近100次
                'performance_metrics': self.performance_metrics
            }
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(state, f, ensure_ascii=False, indent=2)
            
            print(f"✅ 状态已保存到: {filepath}")
            return True
            
        except Exception as e:
            print(f" 保存状态失败: {e}")
            return False
    
    def load_state(self, filepath: str) -> bool:
        """加载管理器状态"""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                state = json.load(f)
            
            self.lottery_type = state.get('lottery_type', self.lottery_type)
            self.config = state.get('config', {})
            self.algorithm_weights = state.get('algorithm_weights', self.default_weights)
            self.prediction_history = state.get('prediction_history', [])
            self.performance_metrics = state.get('performance_metrics', {})
            
            print(f"✅ 状态已从 {filepath} 加载")
            return True
            
        except Exception as e:
            print(f" 加载状态失败: {e}")
            return False


class SimplifiedEnsembleManager:
    """
    简化的集成管理器
    用于在没有复杂依赖的环境中工作
    """
    
    def __init__(self, lottery_type: str = 'ssq'):
        self.lottery_type = lottery_type
        self.algorithms = {}
        self.prediction_history = []
        
        if lottery_type == 'ssq':
            self.red_range = (1, 33)
            self.blue_range = (1, 16) 
            self.red_count = 6
            self.blue_count = 1
        else:  # dlt
            self.red_range = (1, 35)
            self.blue_range = (1, 12)
            self.red_count = 5
            self.blue_count = 2
    
    def register_simple_algorithm(self, name: str, predict_func):
        """注册简单的预测函数"""
        self.algorithms[name] = predict_func
        print(f"✅ 简化算法已注册: {name}")
    
    def generate_simple_predictions(self, num_predictions: int = 3) -> List[Dict]:
        """生成简化的集成预测"""
        import random
        
        ensemble_predictions = []
        
        for i in range(num_predictions):
            # 生成红球
            red_balls = sorted(random.sample(
                range(self.red_range[0], self.red_range[1] + 1),
                self.red_count
            ))
            
            # 生成蓝球
            blue_balls = random.sample(
                range(self.blue_range[0], self.blue_range[1] + 1),
                self.blue_count
            )
            
            ensemble_predictions.append({
                'red': red_balls,
                'blue': blue_balls,
                'algorithm': f'SimplifiedEnsemble-{i+1}',
                'confidence': random.uniform(0.4, 0.8)
            })
        
        return ensemble_predictions